#!/usr/bin/env python
from importlib import import_module
import os, gc, threading, time
import cv2
from flask import Flask, render_template, Response, jsonify
from imageai.Detection import ObjectDetection
import rospy
from std_msgs.msg import String



app = Flask(__name__)

# ROS node, publisher, and parameter.
# The node is started in a separate thread to avoid conflicts with Flask.
# The parameter *disable_signals* must be set if node is not initialized
# in the main thread.

threading.Thread(target=lambda: rospy.init_node('flask_node', disable_signals=True)).start()
pubMotion = rospy.Publisher('/requestMotion01', String, queue_size=1)
pubStop = rospy.Publisher('/requestStop01', String, queue_size=1)




class Camera:

    def __init__(self):
        self.rtsp = "rtsp://admin:js123456@192.168.50.70:554/Streaming/Channels/102"
        self.stack = []
        self.top = 50
        self.cur_frame = None

    def write(self) -> None:
        print('Process to write: %s' % os.getpid())
        cap = cv2.VideoCapture(self.rtsp)
        print("VideoCapture")
        if not cap.isOpened():
            raise RuntimeError('Could not start camera.')
        while True:
            _, img = cap.read()
            if _:
                # print("length {}".format(len(self.stack)))
                self.stack.append(img)
                if len(self.stack) >= self.top:
                    """ 每到一定容量清空一次缓冲栈 """
                    del self.stack[:]
                    gc.collect()
    

    def beginWriter(self) -> None:
        for i in range(1):
            thread_wtite = threading.Thread(target=self.write)
            thread_wtite.start()


    def getFrame(self):
        print("stack的长度", len(self.stack))
        if len(self.stack) <= 0:
            return None
        else:
            value = self.stack.pop()

            if len(self.stack) >= self.top:
                del self.stack[:]
                gc.collect()
            return value


# 构建相机实例对象
myCam = Camera()
execution_path = os.getcwd()

@app.route('/')
def index():
    """Video streaming home page."""
    return render_template('index.html')


def gen(detector):
    """ 图片生成器 """
    while True:
        time.sleep(0.01)
        frame = myCam.getFrame()
        if frame is not None:
            print(type(frame))
            # 执行检测程序
            detections = detector.detectObjectsFromImage(
                                input_image=frame, 
                                minimum_percentage_probability=50,
                                output_type='array')
            
            detec_frame = cv2.imencode('.jpg', detections[0])[1].tobytes()
            yield (b'--frame\r\n' 
                   b'Content-Type: image/jpeg\r\n\r\n' + detec_frame + b'\r\n--frame\r\n')
            


@app.route('/video_feed')
def video_feed():
    myCam.beginWriter()     # 开始往缓冲区写图片
    
    print("execution_path: ", execution_path)
    detector = ObjectDetection()
    detector.setModelTypeAsTinyYOLOv3()
    detector.setModelPath(os.path.join(execution_path, "./src/ros_flask/tiny-yolov3.pt"))
    detector.loadModel()
    return Response(gen(detector), mimetype='multipart/x-mixed-replace; boundary=frame')


if __name__ == '__main__':
    app.run(host='0.0.0.0', threaded=True, use_reloader=False, debug=True)
